"""
车牌识别模块
使用OCR技术识别车牌图像中的文字内容
"""

import re
import cv2
import numpy as np
from typing import List, Dict, Optional, Tuple
from pathlib import Path
import easyocr
from src.config import config
from src.utils.logger import app_logger
from src.utils.image_utils import image_processor


class PlateRecognizer:
    """
    车牌识别器类
    
    用途：
    - 使用OCR技术识别车牌文字
    - 支持中英文车牌识别
    - 提供车牌文字的后处理和验证
    
    使用场景：
    - 从检测到的车牌图像中提取文字
    - 车辆管理系统中的车牌号码识别
    - 交通监控系统中的车牌信息提取
    """
    
    def __init__(self, languages: List[str] = None):
        """
        初始化车牌识别器
        
        参数：
            languages (List[str]): 支持的语言列表，默认使用配置中的语言
        """
        self.languages = languages or config.get('model.ocr_languages')
        self.reader = None
        
        # 中国车牌正则表达式模式（从配置文件读取）
        self.chinese_plate_patterns = config.get('recognition.chinese_plate_patterns')
        
        self._initialize_reader()
    
    def _initialize_reader(self):
        """
        初始化OCR阅读器
        """
        try:
            app_logger.info(f"初始化OCR阅读器，支持语言: {self.languages}")
            
            # 确保模型目录存在
            models_dir = Path(config.get('paths.models_dir'))
            models_dir.mkdir(parents=True, exist_ok=True)
            
            # 创建OCR模型子目录
            ocr_model_dir = models_dir / 'ocr'
            ocr_model_dir.mkdir(parents=True, exist_ok=True)
            
            gpu_enabled = config.get('model.gpu_enabled')
            
            # 指定模型存储目录为项目的models/ocr目录
            app_logger.info(f"OCR模型将存储在: {ocr_model_dir}")
            self.reader = easyocr.Reader(
                self.languages, 
                gpu=gpu_enabled,
                model_storage_directory=str(ocr_model_dir)
            )
            
            app_logger.info("OCR阅读器初始化成功")
            
        except Exception as e:
            app_logger.error(f"初始化OCR阅读器失败: {e}")
            raise
    
    def recognize_plate(self, plate_image: np.ndarray) -> Dict:
        """
        识别单个车牌图像
        
        参数：
            plate_image (np.ndarray): 车牌图像，BGR格式
            
        返回值：
            Dict: 识别结果，包含：
                - text: 识别的文字
                - confidence: 置信度
                - processed_text: 后处理的文字
                - is_valid: 是否为有效车牌
        """
        if self.reader is None:
            app_logger.error("OCR阅读器未初始化")
            return self._empty_result()
        
        try:
            # 预处理图像
            processed_image = image_processor.preprocess_for_ocr(plate_image)
            
            # 使用OCR识别
            results = self.reader.readtext(processed_image)
            
            if not results:
                app_logger.warning("未识别到任何文字")
                return self._empty_result()
            
            # 过滤低置信度结果
            min_confidence = 0.1  # 最低置信度阈值
            filtered_results = [r for r in results if r[2] >= min_confidence]
            
            if not filtered_results:
                app_logger.warning(f"所有识别结果置信度过低 (< {min_confidence})")
                return self._empty_result()
            
            # 获取最佳识别结果
            best_result = self._get_best_ocr_result(filtered_results)
            
            # 后处理文字
            processed_text = self._postprocess_text(best_result['text'])
            
            # 验证车牌格式
            is_valid = self._validate_plate_format(processed_text)
            
            result = {
                'text': best_result['text'],
                'confidence': best_result['confidence'],
                'processed_text': processed_text,
                'is_valid': is_valid
            }
            
            app_logger.info(f"识别结果: {processed_text}, 置信度: {best_result['confidence']:.3f}")
            return result
            
        except Exception as e:
            app_logger.error(f"车牌识别失败: {e}")
            return self._empty_result()
    
    def recognize_plates_batch(self, plate_images: List[np.ndarray]) -> List[Dict]:
        """
        批量识别车牌图像
        
        参数：
            plate_images (List[np.ndarray]): 车牌图像列表
            
        返回值：
            List[Dict]: 识别结果列表
        """
        results = []
        
        for i, plate_image in enumerate(plate_images):
            app_logger.debug(f"识别第 {i+1}/{len(plate_images)} 个车牌")
            result = self.recognize_plate(plate_image)
            results.append(result)
        
        return results
    
    def _get_best_ocr_result(self, ocr_results: List[Tuple]) -> Dict:
        """
        从OCR结果中选择最佳结果
        
        参数：
            ocr_results (List[Tuple]): OCR结果列表
            
        返回值：
            Dict: 最佳结果
        """
        if not ocr_results:
            return {'text': '', 'confidence': 0.0}
        
        # 按置信度排序
        sorted_results = sorted(ocr_results, key=lambda x: x[2], reverse=True)
        
        # 选择置信度最高的结果
        best = sorted_results[0]
        
        return {
            'text': best[1],
            'confidence': best[2]
        }
    
    def _postprocess_text(self, text: str) -> str:
        """
        后处理识别的文字
        
        参数：
            text (str): 原始识别文字
            
        返回值：
            str: 后处理的文字
        """
        if not text:
            return ""
        
        # 移除空格和特殊字符
        processed = re.sub(r'[^\u4e00-\u9fa5A-Z0-9]', '', text.upper())
        
        # 字符替换（常见OCR错误，从配置文件读取）
        char_replacements = config.get('recognition.character_replacements')
        
        # 只在数字位置进行替换
        result = ""
        for i, char in enumerate(processed):
            # 中国车牌格式：第一位是汉字，第二位是字母，后面是字母数字混合
            if i >= 2 and char in char_replacements:
                # 如果周围都是数字，则进行替换
                if i > 0 and i < len(processed) - 1:
                    prev_char = processed[i-1]
                    next_char = processed[i+1]
                    if prev_char.isdigit() or next_char.isdigit():
                        result += char_replacements[char]
                    else:
                        result += char
                else:
                    result += char
            else:
                result += char
        
        return result
    
    def _validate_plate_format(self, text: str) -> bool:
        """
        验证车牌格式是否有效
        
        参数：
            text (str): 车牌文字
            
        返回值：
            bool: 是否为有效车牌格式
        """
        if not text or len(text) < 6:
            return False
        
        # 检查是否匹配中国车牌格式
        for pattern in self.chinese_plate_patterns:
            if re.match(pattern, text):
                return True
        
        # 检查基本格式（至少包含字母和数字）
        has_letter = any(c.isalpha() for c in text)
        has_digit = any(c.isdigit() for c in text)
        
        return has_letter and has_digit
    
    def _empty_result(self) -> Dict:
        """
        返回空的识别结果
        
        返回值：
            Dict: 空结果
        """
        return {
            'text': '',
            'confidence': 0.0,
            'processed_text': '',
            'is_valid': False
        }
    
    def enhance_plate_image(self, plate_image: np.ndarray) -> np.ndarray:
        """
        增强车牌图像以提高识别准确率
        
        参数：
            plate_image (np.ndarray): 车牌图像
            
        返回值：
            np.ndarray: 增强后的车牌图像
        """
        # 转换为灰度图
        if len(plate_image.shape) == 3:
            gray = cv2.cvtColor(plate_image, cv2.COLOR_BGR2GRAY)
        else:
            gray = plate_image.copy()
        
        # 直方图均衡化
        equalized = cv2.equalizeHist(gray)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(equalized, (3, 3), 0)
        
        # 锐化
        kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
        sharpened = cv2.filter2D(blurred, -1, kernel)
        
        return sharpened
    
    def get_character_boxes(self, plate_image: np.ndarray) -> List[Dict]:
        """
        获取车牌中每个字符的位置框
        
        参数：
            plate_image (np.ndarray): 车牌图像
            
        返回值：
            List[Dict]: 字符框列表，每个包含位置和文字
        """
        if self.reader is None:
            return []
        
        try:
            processed_image = image_processor.preprocess_for_ocr(plate_image)
            results = self.reader.readtext(processed_image, detail=1)
            
            character_boxes = []
            for result in results:
                bbox, text, confidence = result
                
                # 转换边界框格式
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                
                char_box = {
                    'bbox': (min(x_coords), min(y_coords), max(x_coords), max(y_coords)),
                    'text': text,
                    'confidence': confidence
                }
                
                character_boxes.append(char_box)
            
            return character_boxes
            
        except Exception as e:
            app_logger.error(f"获取字符框失败: {e}")
            return []
    
    def set_languages(self, languages: List[str]):
        """
        设置支持的语言
        
        参数：
            languages (List[str]): 语言列表
        """
        self.languages = languages
        self._initialize_reader()
        app_logger.info(f"语言设置更新为: {languages}")


# 全局车牌识别器实例
plate_recognizer = PlateRecognizer() 